Writing a faster mean-squared-displacement function

I'm trying to calculate the Mean Squared Displacement (MSD) of a particle trajectory. In reality, I have an array of positions for numtrj particles and numpts timepoints and dim dimensions: pos[numtrj, numpts, dim]. I think my question has the same answer if I just have the x trajectory of a single particle, though.

In case you haven't done MSD calculations before, there's one cute way in which you get extra information out of a trajectory. Say I have just five time points, and my positions are

In [83]: x

Out[83]: array([ 0.        ,  1.74528704,  1.59639865,  2.59976219,  3.70852457])

You could just get squared displacement by looking at x**2. However, you could also say that you have 4 different values for the displacement at one timestep:

x[1] - x[0], x[2] - x[1], x[3] - x[2], x[4] - x[3]

Similarly, three values for displacement at two timesteps: x[2:] - x[:-2]. Etc. So, the way I'm calculating MSD at the moment is:

def msd_1d(x):
    result = np.zeros_like(x)
    for i in range(1,len(x)):
        result[i] = np.average((x[i:] - x[:-i])**2)
    return result

or

def get_msd_traj(pos):
    result = np.zeros_like(pos)
    for i in range(1,pos.shape[1]):
        result[:,i,:] = np.average((pos[:,i:,:] - pos[:,:-i,:])**2,axis=1)
    return result

(side note: often the data comes indexed like pos[numpts, numtrj, dim] for molecular dynamics trajectories, but that doesn't change anything here)

So, I asked Joshua Adelman if he had any quick thoughts.

Read more…